# Copyright 2021 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
import six
import argparse


def ubuntu_parse_args():
    """
    Deep Attention Matching Network Config
    """
    parser = argparse.ArgumentParser("DAM_ubuntu Config")

    parser.add_argument('--model_name', type=str, default="DAM_ubuntu", help=' ')
    parser.add_argument('--device_target', type=str, default="Ascend", choices=("CPU", "GPU", "Ascend"),
                        help="run platform, only support CPU, GPU and Ascend", )
    parser.add_argument('--do_eval', type=bool, default=True,
                        help="Whether side training changes verification.")
    parser.add_argument('--is_emb_init', type=bool, default=True, help='')

    # 文件地址
    parser.add_argument('--data_path', type=str, default="./data/ubuntu/data.pkl",
                        help='Path to training data. (default: %(default)s)')
    parser.add_argument('--train_data_path', type=str, default="./data/ubuntu/data_train.mindrecord",
                        help='Path to training data. (default: %(default)s)')
    parser.add_argument('--eval_data_path', type=str, default="./data/ubuntu/data_val.mindrecord",
                        help='Path to eval data. (default: %(default)s)')
    parser.add_argument('--test_data_path', type=str, default="./data/ubuntu/data_test.mindrecord",
                        help='Path to testing data. (default: %(default)s)')
    parser.add_argument('--save_path', type=str, default="./checkpoints/ubuntu/temp/",
                        help='Path to save trained models. (default: %(default)s)')
    parser.add_argument('--emb_init', type=str, default='./data/ubuntu/word_embedding.pkl')
    parser.add_argument('--model_path', type=str, default="./checkpoints/ubuntu/temp/",
                        help='Path to load well-trained models. (default: %(default)s)')
    parser.add_argument('--output_path', type=str, default="./output/ubuntu/", help='')
    parser.add_argument('--eval_file_name', type=str, default="./log/eval_accuracy_ubuntu",
                        help='')
    parser.add_argument('--loss_file_name', type=str, default="./log/loss_ubuntu",
                        help='Loss log file path. Default: "./loss_ubuntu.log"')
    parser.add_argument('--pretrained_model', type=str, default=None,
                        help='the pre-trained file of model.')
    parser.add_argument('--checkpoint_path', type=str, default=None, help="checkpoint root path.")
    parser.add_argument('--checkpoint_name', type=str, default=None, help="eval checkpoint.")

    # 数据格式、网络结构参数
    parser.add_argument('--max_turn_num', type=int, default=9, help='Maximum number of utterances in context.')
    parser.add_argument('--max_turn_len', type=int, default=50,
                        help='Maximum length of setences in turns.')  # turns 应该是每一轮对话中的正文
    parser.add_argument('--vocab_size', type=int, default=434512, help='The size of vocabulary.')
    # 434512 for ubuntu data; 172131 for douban data
    parser.add_argument('--emb_size', type=int, default=200, help='The dimension of word embedding')
    parser.add_argument('--_EOS_', type=int, default=28270, help='The id for the end of sentence in vocabulary.')
    # 1 for douban data; 28270 for ubuntu data
    parser.add_argument('--stack_num', type=int, default=5, help='The number of stacked attentive modules in network.')
    parser.add_argument('--channel1_dim', type=int, default=32,
                        help='The channels number of the 1st conv3d layers output.')
    parser.add_argument('--channel2_dim', type=int, default=16,
                        help='The channels number of the 2nd conv3d layers output.')
    parser.add_argument('--is_mask', type=bool, default=True, help='')
    parser.add_argument('--is_layer_norm', type=bool, default=True, help='')
    parser.add_argument('--is_positional', type=bool, default=False, help='')

    # trick
    parser.add_argument('--batch_size', type=int, default=256, help='Batch size for training. (default: %(default)d)')
    parser.add_argument('--eval_batch_size', type=int, default=200,
                        help='Batch size for training. (default: %(default)d)')
    parser.add_argument('--learning_rate', type=float, default=1e-3,
                        help='Learning rate used to train. (default: %(default)f)')
    parser.add_argument('--decay_rate', type=float, default=0.9, help='The decay rate.')
    parser.add_argument('--decay_steps', type=int, default=400,
                        help='A value used to calculate decayed learning rate.')
    parser.add_argument('--loss_scale', type=int, default=1)
    parser.add_argument('--load_pretrained', type=bool, default=False)
    parser.add_argument('--epoch_size', type=int, default=2)
    parser.add_argument('--drop_prob', type=float, default=None)

    parser.add_argument('--data_url', required=True, default="obs://harbin-engineering-uni/DAM/data/ubuntu/",
                        help='Location of data.')
    parser.add_argument('--train_url', required=True, default="obs://harbin-engineering-uni/DAM/output/ubuntu/",
                        help='Location of training outputs.')
    parser.add_argument('--time_', type=int, default=None, help="version")

    args = parser.parse_args()
    return args


def douban_parse_args():
    """
    Deep Attention Matching Network Config
    """
    parser = argparse.ArgumentParser("DAM_douban Config")

    parser.add_argument('--model_name', type=str, default="DAM_douban", help='')
    parser.add_argument('--device_target', type=str, default="Ascend", choices=("CPU", "GPU", "Ascend"),
                        help="run platform, only support CPU, GPU and Ascend", )
    parser.add_argument('--do_eval', type=bool, default=True,
                        help="Whether side training changes verification.")
    parser.add_argument('--is_emb_init', type=bool, default=True, help='')
    # 文件地址
    parser.add_argument('--data_path', type=str, default="./data/douban/data.pkl",
                        help='Path to training data. (default: %(default)s)')
    parser.add_argument('--train_data_path', type=str, default="./data/douban/data_train.mindrecord",
                        help='Path to training data. (default: %(default)s)')
    parser.add_argument('--eval_data_path', type=str, default="./data/douban/data_test.mindrecord",
                        help='Path to eval data. (default: %(default)s)')
    parser.add_argument('--test_data_path', type=str, default="./data/douban/data_test.mindrecord",
                        help='Path to testing data. (default: %(default)s)')
    parser.add_argument('--save_path', type=str, default="./checkpoints/douban/temp/",
                        help='Path to save trained models. (default: %(default)s)')
    parser.add_argument('--emb_init', type=str, default="./data/douban/word_embedding.pkl")
    parser.add_argument('--model_path', type=str, default="./checkpoints/douban/temp/",
                        help='Path to load well-trained models. (default: %(default)s)')
    parser.add_argument('--output_path', type=str, default="./output/douban/")
    parser.add_argument('--eval_file_name', type=str, default="./log/eval_accuracy_douban")
    parser.add_argument('--loss_file_name', type=str, default="./log/loss_douban",
                        help='Loss log file path. Default: "./log/loss_douban.log"')
    parser.add_argument('--pretrained_model', type=str, default=None,
                        help='the pre-trained file of model.')
    parser.add_argument('--checkpoint_path', type=str, default=None, help="checkpoint root path.")
    parser.add_argument('--checkpoint_name', type=str, default=None, help="eval checkpoint.")

    # 数据格式、网络结构参数
    parser.add_argument('--max_turn_num', type=int, default=9, help='Maximum number of utterances in context.')
    parser.add_argument('--max_turn_len', type=int, default=50,
                        help='Maximum length of setences in turns.')  # turns 应该是每一轮对话中的正文
    parser.add_argument('--vocab_size', type=int, default=172130, help='The size of vocabulary.')
    # 434512 for ubuntu data; 172130 for douban data
    parser.add_argument('--emb_size', type=int, default=200, help='The dimension of word embedding')
    parser.add_argument('--_EOS_', type=int, default=1, help='The id for the end of sentence in vocabulary.')
    # 1 for douban data; 28270 for ubuntu data
    parser.add_argument('--stack_num', type=int, default=5, help='The number of stacked attentive modules in network.')
    parser.add_argument('--channel1_dim', type=int, default=16,
                        help='The channels number of the 1st conv3d layers output.')
    parser.add_argument('--channel2_dim', type=int, default=16,
                        help='The channels number of the 2nd conv3d layers output.')
    parser.add_argument('--is_mask', type=bool, default=True, help='')
    parser.add_argument('--is_layer_norm', type=bool, default=True, help='')
    parser.add_argument('--is_positional', type=bool, default=False, help='')

    # trick
    parser.add_argument('--batch_size', type=int, default=256, help='Batch size for training. (default: %(default)d)')
    parser.add_argument('--eval_batch_size', type=int, default=256,
                        help='Batch size for training. (default: %(default)d)')
    parser.add_argument('--learning_rate', type=float, default=1e-3,
                        help='Learning rate used to train. (default: %(default)f)')  # 1e-3
    parser.add_argument('--decay_rate', type=float, default=0.9, help='The decay rate.')
    parser.add_argument('--decay_steps', type=int, default=400,
                        help='A value used to calculate decayed learning rate.')
    parser.add_argument('--loss_scale', type=int, default=1, help='')
    parser.add_argument('--load_pretrained', type=bool, default=False, help='')
    parser.add_argument('--epoch_size', type=int, default=2)
    parser.add_argument('--drop_prob', type=float, default=None)

    parser.add_argument('--data_url', required=True, default="obs://harbin-engineering-uni/DAM/data/ubuntu/",
                        help='Location of data.')
    parser.add_argument('--train_url', required=True, default="obs://harbin-engineering-uni/DAM/output/ubuntu/",
                        help='Location of training outputs.')
    parser.add_argument('--time_', type=int, default=None, help="version")

    args = parser.parse_args()
    return args


def print_arguments(args):
    """
    Print Config
    """
    print('-----------  Configuration Arguments -----------')
    for arg, value in sorted(six.iteritems(vars(args))):
        print('%s: %s' % (arg, value))
    print('------------------------------------------------')
